import cv2
import numpy as np


class Otsu:
    def __init__(self, type):
        self.type = type
    
    def __call__(self, src):
        """
            OTSU segmenting
            type: type of cost function
        """

        h, w = src.shape   # height and weight

        dst = np.zeros_like(src).tolist() # zero

        src = src.tolist() # list Object is faster than numpy Object

        score = [0 for _ in range(0, 256)] # score for each grayscale

        histogram = [0 for _ in range(0, 256)] # histogram

        # histogram caculate
        for y in range(0, h):
            for x in range(0, w):
                histogram[src[y][x]] += 1
        
        # grayscale range
        min_t = 255
        for t in range(0, 256):
            if histogram[t]:
                min_t = t
                break
       
        # grayscale range
        max_t = 0
        for t in range(255, -1, -1):
            if histogram[t]:
                max_t = t
                break

        # histogram nomaliza
        s = h * w
        for i in range(0, 256):
            histogram[i] /= s


        # weights 
        weights = [0 for _ in range(0, 256)]
        weights[0] = histogram[0]
        for i in range(1, 256):
            weights[i] = weights[i - 1] + histogram[i]

        # mean
        mean = 0
        for i in range(0, 256):
            mean += histogram[i] * i

        if self.type == 'INTER_MIN': # 最小类间差
            def evaluate(t, fmean, bmean):
                fvar = 0
                for i in range(0, t):
                    fvar += histogram[i] * ((i - fmean) ** 2)
                bvar = 0
                for i in range(t, 256):
                    bvar += histogram[i] * ((i - bmean) ** 2)
                return fvar + bvar
            def compare(a, b):
                if a < b:
                    return True
        elif self.type == 'INNER_MAX': # 最大类内差
            def evaluate(t, fmean, bmean):
                return weights[t] * (1 - weights[t]) * ((bmean - fmean) ** 2)
                # return weights[t] * ((mean - fmean) ** 2) + (1 - weights[t]) * ((mean - bmean) ** 2)
            def compare(a, b):
                if a > b:
                    return True

        # score
        score = [0 for _ in range(0, 256)]
        f = 0
        b = mean
        for t in range(min_t, max_t):
            f += histogram[t] * t
            b -= histogram[t] * t
            score[t] = evaluate(t, f / weights[t], b / (1 - weights[t]))

        # best threshold
        best_score = score[min_t]
        best_t = min_t
        for t in range(min_t + 1, max_t):
            if compare(score[t], best_score):
                best_score = score[t]
                best_t = t

        # threshold table
        threstable = [0 for _ in range(0, best_t + 1)] + [255 for _ in range(best_t + 1, 256)]

        # threshold
        for y in range(0, h):
            for x in range(0, w):
                dst[y][x] = threstable[src[y][x]]
        
        dst = np.array(dst).astype('uint8')

        return best_t, dst


if __name__ == '__main__':

    SRC = cv2.imread('lena.jpg', cv2.IMREAD_GRAYSCALE)

    def show(winname, img):
        cv2.imshow(winname, img)
        cv2.resizeWindow(winname, 512, 512)
        cv2.waitKey(1)

    t, dst = Otsu('INTER_MIN')(SRC)
    show('INTER_MIN', dst)

    t, dst = Otsu('INNER_MAX')(SRC)
    show('INNER_MAX', dst)

    t, dst = cv2.threshold(SRC, 0, 255, cv2.THRESH_OTSU)
    show('OPENCV_OTSU', dst)

    cv2.waitKey(0)